Airflowのユニットテストモード実施を嵌りなく進めるための設定をまとめてみた

はじめに

Airflowの機能を一部使ったロジックのテストを検討しました。問題は既にあるユニットテストに載せられるのかどうか。

テストを実行するまで必要な手続きをまとめてみました。

Airflowが読み込む設定ファイル

Airflowの設定は、以下の順で優先されます。

1. AIRFLOW__<SECRION>__<KEY>に設定された値
2. AIRFLOW_CONFIGに設定されたファイルの内容
3. airflow.cfg

では、テスト用の構成はどうなるか、というと

1. AIRFLOW__<SECRION>__<KEY>に設定された値
2. AIRFLOW_TEST_CONFIGに設定されたファイル
3. unittests.cfg

の順になります。

Airflowをユニットテストモードにする

では、テスト設定にしたい場合に必要なものは何かというと、unit_test_modeのフラグです。ファイルには以下のように設定します。

[core]
unit_test_mode = True

ただし、単純にunittests.cfgに書くだけでは認識されません。何故なら、unit_test_modeのフラグがTrueとなっていて初めて読み込まれるファイルだからです。とはいえairflow.cfgに書いてしまうと常時テストモードになってしまいます。

そこで、以下の設定を順に行います。なお、AIRFLOW_HOMEについては設定済みの前提としています。

環境変数でUNIT_TEST_MODEに指定する

以下のように環境変数へ指定します。COREは両側を__(アンダーバーx2)で囲む点だけ注意です。

export AIRFLOW__CORE__UNIT_TEST_MODE=True

環境変数でAIRFLOW_CONFIGにunittests.cfgを指定する

テスト用設定ファイルを指定します。標準指定のファイル名として敢えて指定する理由があります。

UNIT_TEST_MODEへの指定のみでもテストモードにはなるのですが、設定内のsql_alchemy_connの初期値が原因でvalidateエラーを発生させる可能性があります。

..: in <module>
    from airflow.operators.dagrun_operator import DagRunOrder
.venv/lib/python3.7/site-packages/airflow/__init__.py:31: in <module>
    from airflow.utils.log.logging_mixin import LoggingMixin
.venv/lib/python3.7/site-packages/airflow/utils/__init__.py:24: in <module>
    from .decorators import apply_defaults as _apply_defaults
.venv/lib/python3.7/site-packages/airflow/utils/decorators.py:36: in <module>
    from airflow import settings
.venv/lib/python3.7/site-packages/airflow/settings.py:38: in <module>
    from airflow.configuration import conf, AIRFLOW_HOME, WEBSERVER_CONFIG  # NOQA F401
.venv/lib/python3.7/site-packages/airflow/configuration.py:731: in <module>
    conf.read(AIRFLOW_CONFIG)
.venv/lib/python3.7/site-packages/airflow/configuration.py:421: in read
    self._validate()
.venv/lib/python3.7/site-packages/airflow/configuration.py:213: in _validate
    self._validate_config_dependencies()
.venv/lib/python3.7/site-packages/airflow/configuration.py:247: in _validate_config_dependencies
    self.get('core', 'executor')))
E   airflow.exceptions.AirflowConfigException: error: cannot use sqlite with the CeleryExecutor

標準設定指定のファイルの中身を読み込んだ上で、上書きする形でテスト設定指定のファイルが読み込まれます。本番と設定が異なる可能性もありますし、ダミーのファイル設置は無駄な管理コストが増えてしまいます。設定を書き換えれば問題ないので、unittests.cfgを以下のように変更した上で、標準設定ファイルとして扱ってしまいましょう。

sql_alchemy_conn = mysql://root:000000@localhost/airflow

export AIRFLOW_CONFIG=unittests.cfg

必要に応じて、mysqlclientをインストールしてください。

pipenv install mysqlclient --dev

pytestで手軽に扱う

いずれも、pytestを使っている場合はpytest.iniのenvセクションに書いておくと、コマンド実行時に環境変数への指定の手間が省けます。

env = 
      AIRFLOW__CORE__UNIT_TEST_MODE = True
      AIRFLOW_CONFIG = unittests.cfg

あとがき

ユニットテスト用のフラグだけ立てても上手くいかず、テスト設定用ファイルだけを書き換えても反応せず、結果設定読み込み用クラスを辿ることになりました。

動作を理解するまでは嵌り方が酷いことになると思いますが、分かってしまえば管理するべき対象も、要件にそったカスタマイズも容易になると思います。困ったときの参考になれば幸いです。